import os
import requests
from transformers import AutoTokenizer, AutoModelForCausalLM

# 设置环境变量解决SSL问题
os.environ['CURL_CA_BUNDLE'] = ''
# 替换为清华源，实际测试因为内存太大，每次下载需要特别长时间，想下载完成需要多次尝试
os.environ['HF_ENDPOINT'] = 'https://mirrors.tuna.tsinghua.edu.cn/hugging-face-models'


# 模型比较大，12G，需要修改地址
# 设置模型缓存目录（建议放在D盘或其他空间较大的盘）
cache_dir = "D:/demo/gitee/python/models/chatglm3-6b"
os.makedirs(cache_dir, exist_ok=True)

# 增加请求的超时时间
requests.adapters.DEFAULT_RETRIES = 5
session = requests.Session()

session.keep_alive = False


model_name = "THUDM/chatglm3-6b"
print("开始下载模型...")

# 下载设置
kwargs = {
    "trust_remote_code": True,
    "cache_dir": cache_dir,
    "local_files_only": False,
    "resume_download": True,  # 启用断点续传
    "timeout": 1000,         # 增加超时时间到1000秒
    "max_retries": 10        # 增加重试次数
}


# 下载tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs)
# tokenizer = AutoTokenizer.from_pretrained(
#     model_name, 
#     trust_remote_code=True,
#     cache_dir=cache_dir,
#     timeout=1000,  # 增加超时时间到600秒
#     max_retries=10  # 增加重试次数
# )

# 下载模型
model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
# model = AutoModelForCausalLM.from_pretrained(
#     model_name, 
#     trust_remote_code=True,
#     cache_dir=cache_dir,
#     timeout=1000,
#     max_retries=10
# )

print("模型下载完成")